import pandas as pd
import numpy as np
from func import close_and_mcap,get_daily_r,dump_to_db,dump_cash,buy_stock,drop_stock,crt1,crt2,crt3




ori_data = pd.read_excel('stock_pool.xlsx')
# ori_data已经按r降序排列
init_amt = 10**(9) #初始规模
start_date = '2013-01-04'
update_date_list = [start_date]

sharps = []
result = {}


for N in range(35,51):
    # 为满足地域和行业要求，N至少为35
    data = ori_data.iloc[:N,:]
    # data已经按r降序排列

    code = tuple(data.fic_code)

    close_mcap = pd.DataFrame(close_and_mcap('2013-01-04', code).to_dict())
    r_close = data.merge(close_mcap, on='fic_code')[['fic_code', 'r', 'close']]

    shares_change = buy_stock(r_close, init_amt)

    shares_close_mcap = shares_change.merge(close_mcap, on='fic_code')
    shares_close_mcap.columns = ['fic_code', 'shares', 'close', 'mktcap']
    cash = init_amt - (shares_close_mcap['close'] * shares_close_mcap['shares']).sum()

    while len(crt2(shares_close_mcap)) > 0 or len(crt3(shares_close_mcap)) > 0 or crt1(shares_close_mcap, cash):
        stock_need_drop = set((crt2(shares_close_mcap))) | set(crt3(shares_close_mcap))
        print('不满足条件，需要减仓的股票是:{}'.format(stock_need_drop))
        shares_change = drop_stock(r_close, shares_close_mcap[['fic_code', 'shares']], 0, 0, 1, stock_need_drop)
        shares_close_mcap = shares_close_mcap.merge(shares_change,on='fic_code')
        shares_close_mcap['shares'] = shares_close_mcap['shares'] - shares_close_mcap['shares_change']
        shares_close_mcap = shares_close_mcap.iloc[:,:-1] # 把shares_change这列删除
        cash = init_amt - (shares_close_mcap['close'] * shares_close_mcap['shares']).sum()

        # 现金增加，购买股票
        shares_change = buy_stock(r_close,cash)
        shares_close_mcap = shares_close_mcap.merge(shares_change,on='fic_code')
        shares_close_mcap['shares'] = shares_close_mcap['shares'] + shares_close_mcap['shares_change']
        shares_close_mcap = shares_close_mcap.iloc[:,:-1] # 把shares_change这列删除
        cash = init_amt - (shares_close_mcap['close'] * shares_close_mcap['shares']).sum()
        crt1_result = crt1(shares_close_mcap, cash)

    else:
        print('符合要求')

    print('建仓成功，接下来计算sharp')

    daily_r = pd.DataFrame(get_daily_r(code).to_dict())
    daily_avg_r = daily_r.groupby('date').mean()['r']
    start_day_r = daily_r.loc[[str(i) == '2013-01-04' for i in daily_r.date], ['fic_code', 'r']]
    start_day_avg_r = shares_close_mcap.merge(start_day_r, on='fic_code')[['fic_code', 'r', 'shares', 'close']]
    start_day_avg_r['weight'] = start_day_avg_r['shares'] * start_day_avg_r['close'] / np.sum(
        start_day_avg_r['shares'] * start_day_avg_r['close'])
    start_day_avg_r = np.sum(start_day_avg_r['r'] * start_day_avg_r['weight'])
    daily_avg_r[-1] = start_day_avg_r

    volatility = (np.sum(np.power(daily_avg_r - daily_avg_r.mean(), 2)) / (len(daily_avg_r) - 1)) ** (0.5)
    sharp = (daily_avg_r.mean()) / volatility

    print('N为{}时，sharp值为:{}'.format(N,sharp))
    sharps.append((N,sharp))
    result[N] = shares_close_mcap[['fic_code','shares','close']]


sharps = np.array(sharps)
best_idx = sharps.argmax()-30
best_N = range(35,51)[best_idx]
dump_info = result[best_N]
dump_info['security_mkt_value'] = dump_info['shares']*dump_info['close']
dump_info['weight'] = dump_info.security_mkt_value/np.sum(dump_info.security_mkt_value)
dump_info['portfolio_id'] = 110001
dump_info['date'] = start_date
dump_info['cash_flow'] = dump_info['security_mkt_value']
dump_info['value_add'] = 0

print(dump_info)
cash = init_amt-dump_info.security_mkt_value.sum()
print('建仓后还剩:{}'.format(init_amt-dump_info.security_mkt_value.sum()))
dump_info = dump_info.to_dict('records')
dump_to_db(dump_info)
dump_cash(start_date,cash)
